#!/usr/bin/env python3
"""
GPU支持诊断脚本
检查系统GPU状态并提供解决方案
"""

import torch
import sys
import subprocess
import platform

def check_system_info():
    """检查系统基本信息"""
    print("="*60)
    print("系统信息诊断")
    print("="*60)
    
    print(f"操作系统: {platform.system()} {platform.release()}")
    print(f"Python版本: {sys.version}")
    print(f"PyTorch版本: {torch.__version__}")
    
    # 检查CUDA版本
    try:
        cuda_version = torch.version.cuda
        print(f"CUDA版本: {cuda_version if cuda_version else '未安装CUDA版本的PyTorch'}")
    except:
        print("CUDA版本: 无法获取")

def check_gpu_hardware():
    """检查GPU硬件"""
    print("\n" + "="*60)
    print("GPU硬件检查")
    print("="*60)
    
    try:
        if platform.system() == "Windows":
            # Windows系统使用wmic命令
            result = subprocess.run(['wmic', 'path', 'win32_VideoController', 'get', 'name'], 
                                 capture_output=True, text=True, shell=True)
            if result.returncode == 0:
                gpu_info = result.stdout.strip().split('\n')[1:]
                print("检测到的显卡:")
                for gpu in gpu_info:
                    gpu = gpu.strip()
                    if gpu and gpu != "Name":
                        print(f"  - {gpu}")
                        if "NVIDIA" in gpu.upper():
                            print("    ✅ 检测到NVIDIA显卡，支持CUDA")
                        elif "AMD" in gpu.upper() or "RADEON" in gpu.upper():
                            print("    ⚠️ AMD显卡不支持CUDA，但可能支持ROCm")
                        elif "INTEL" in gpu.upper():
                            print("    ℹ️ Intel集成显卡，不支持CUDA")
            else:
                print("无法获取GPU信息")
        else:
            # Linux/Mac系统
            try:
                result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
                if result.returncode == 0:
                    print("nvidia-smi输出:")
                    print(result.stdout)
                else:
                    print("nvidia-smi命令失败，可能没有安装NVIDIA驱动")
            except FileNotFoundError:
                print("未找到nvidia-smi命令")
                
    except Exception as e:
        print(f"GPU硬件检查失败: {e}")

def check_pytorch_cuda():
    """检查PyTorch的CUDA支持"""
    print("\n" + "="*60)
    print("PyTorch CUDA支持检查")
    print("="*60)
    
    print(f"torch.cuda.is_available(): {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"✅ CUDA可用!")
        print(f"CUDA设备数量: {torch.cuda.device_count()}")
        
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}:")
            print(f"  名称: {torch.cuda.get_device_name(i)}")
            print(f"  计算能力: {torch.cuda.get_device_capability(i)}")
            
            # 获取显存信息
            torch.cuda.set_device(i)
            total_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
            print(f"  总显存: {total_memory:.1f} GB")
            
            # 检查当前显存使用
            if torch.cuda.is_available():
                allocated = torch.cuda.memory_allocated(i) / 1024**3
                cached = torch.cuda.memory_reserved(i) / 1024**3
                print(f"  已分配显存: {allocated:.1f} GB")
                print(f"  缓存显存: {cached:.1f} GB")
    else:
        print("❌ CUDA不可用")
        print("\n可能的原因:")
        print("1. 没有NVIDIA GPU")
        print("2. 安装了CPU版本的PyTorch")
        print("3. CUDA驱动版本不兼容")
        print("4. GPU被禁用")

def test_gpu_performance():
    """测试GPU性能"""
    if not torch.cuda.is_available():
        print("\n跳过GPU性能测试（CUDA不可用）")
        return
    
    print("\n" + "="*60)
    print("GPU性能测试")
    print("="*60)
    
    try:
        # 简单的GPU计算测试
        device = torch.device('cuda:0')
        
        # 创建测试张量
        print("创建测试张量...")
        x = torch.randn(1000, 1000, device=device)
        y = torch.randn(1000, 1000, device=device)
        
        # 计时测试
        import time
        torch.cuda.synchronize()
        start_time = time.time()
        
        # 执行矩阵乘法
        for _ in range(100):
            z = torch.matmul(x, y)
        
        torch.cuda.synchronize()
        end_time = time.time()
        
        print(f"✅ GPU计算测试成功!")
        print(f"100次矩阵乘法用时: {end_time - start_time:.3f}秒")
        
        # 显存使用测试
        allocated = torch.cuda.memory_allocated(0) / 1024**2
        print(f"当前显存使用: {allocated:.1f} MB")
        
    except Exception as e:
        print(f"❌ GPU性能测试失败: {e}")

def provide_solutions():
    """提供解决方案"""
    print("\n" + "="*60)
    print("解决方案建议")
    print("="*60)
    
    if not torch.cuda.is_available():
        print("🔧 CUDA不可用的解决方案:")
        print("\n1. 检查是否有NVIDIA显卡:")
        print("   - 确认你的笔记本有独立的NVIDIA显卡")
        print("   - 某些笔记本只有Intel集成显卡或AMD显卡")
        
        print("\n2. 安装正确版本的PyTorch:")
        print("   - 卸载当前PyTorch: pip uninstall torch torchvision torchaudio")
        print("   - 访问 https://pytorch.org/get-started/locally/")
        print("   - 选择适合你的CUDA版本")
        print("   - 例如CUDA 11.8: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118")
        print("   - 例如CUDA 12.1: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")
        
        print("\n3. 更新NVIDIA驱动:")
        print("   - 访问NVIDIA官网下载最新驱动")
        print("   - 或使用GeForce Experience自动更新")
        
        print("\n4. 笔记本电源设置:")
        print("   - 确保使用高性能模式")
        print("   - 在NVIDIA控制面板中设置首选GPU")
        
        print("\n5. 检查Windows设置 (如果是Windows):")
        print("   - 设置 > 系统 > 显示 > 图形设置")
        print("   - 为Python选择高性能GPU")
    
    else:
        print("✅ GPU工作正常!")
        print("\n优化建议:")
        print("- 确保电源模式设置为高性能")
        print("- 关闭不必要的后台程序释放显存")
        print("- 根据显存大小调整batch_size")

def generate_quick_train_gpu_fix():
    """生成修改后的quick_train脚本"""
    print("\n" + "="*60)
    print("生成GPU优化的quick_train脚本")
    print("="*60)
    
    fixed_code = '''#!/usr/bin/env python3
"""
GPU优化的快速训练安全帽检测模型
包含详细的GPU检测和错误处理
"""

from ultralytics import YOLO
import torch
import os

def detailed_gpu_check():
    """详细的GPU检查"""
    print("=" * 50)
    print("GPU检测详情")
    print("=" * 50)
    
    print(f"PyTorch版本: {torch.__version__}")
    print(f"CUDA是否可用: {torch.cuda.is_available()}")
    
    if torch.cuda.is_available():
        print(f"CUDA版本: {torch.version.cuda}")
        print(f"GPU数量: {torch.cuda.device_count()}")
        
        for i in range(torch.cuda.device_count()):
            props = torch.cuda.get_device_properties(i)
            print(f"GPU {i}: {props.name}")
            print(f"  显存: {props.total_memory / 1024**3:.1f} GB")
            print(f"  计算能力: {props.major}.{props.minor}")
        
        # 测试GPU是否真正可用
        try:
            test_tensor = torch.tensor([1.0]).cuda()
            print("✅ GPU测试成功")
            return True
        except Exception as e:
            print(f"❌ GPU测试失败: {e}")
            return False
    else:
        print("❌ CUDA不可用")
        print("可能原因:")
        print("1. 没有NVIDIA GPU")
        print("2. 安装了CPU版本的PyTorch")
        print("3. CUDA驱动问题")
        return False

def main():
    print("=" * 50)
    print("GPU优化的快速训练安全帽检测模型")
    print("=" * 50)
    
    # 详细检查GPU
    gpu_available = detailed_gpu_check()
    
    # 强制指定设备
    if gpu_available:
        device = 'cuda:0'  # 明确指定第一个GPU
        print(f"\\n🚀 使用GPU训练: {device}")
        
        # 设置CUDA相关环境变量（可选）
        os.environ['CUDA_LAUNCH_BLOCKING'] = '1'  # 同步CUDA操作，便于调试
        
    else:
        device = 'cpu'
        print(f"\\n⚠️ 回退到CPU训练: {device}")
        print("建议: 运行 python check_gpu_support.py 诊断GPU问题")
    
    # 显示当前设备信息
    if device == 'cuda:0':
        current_device = torch.cuda.current_device()
        device_name = torch.cuda.get_device_name(current_device)
        print(f"当前设备: {device_name}")
    
    # 加载预训练模型
    print("\\n加载YOLOv8预训练模型...")
    model = YOLO('yolov8n.pt')  # nano版本，训练速度快
    
    # 训练参数调整
    if device == 'cuda:0':
        # GPU训练参数
        batch_size = 16  # GPU可以用更大的batch
        workers = 4      # 更多工作线程
    else:
        # CPU训练参数
        batch_size = 8   # CPU用较小batch
        workers = 2      # 较少工作线程
    
    print(f"\\n训练配置:")
    print(f"设备: {device}")
    print(f"批次大小: {batch_size}")
    print(f"工作线程: {workers}")
    
    # 开始训练
    print("\\n开始训练...")
    try:
        results = model.train(
            data='data.yaml',           # 数据集配置文件
            epochs=50,                  # 训练轮数
            imgsz=640,                  # 输入图像大小
            batch=batch_size,           # 批次大小
            device=device,              # 明确指定设备
            workers=workers,            # 数据加载线程数
            project='runs/detect',      # 项目保存目录
            name='helmet_detection',    # 实验名称
            save=True,                  # 保存检查点
            save_period=10,             # 每10轮保存一次
            patience=15,                # 早停耐心值
            cache=False,                # 不缓存图像到内存
            augment=True,               # 使用数据增强
            verbose=True                # 详细输出
        )
        
        print("\\n训练完成!")
        print(f"最佳模型保存在: runs/detect/helmet_detection/weights/best.pt")
        print(f"最新模型保存在: runs/detect/helmet_detection/weights/last.pt")
        
        # 验证模型
        print("\\n验证模型...")
        metrics = model.val(device=device)  # 验证时也使用相同设备
        print(f"mAP50: {metrics.box.map50:.3f}")
        print(f"mAP50-95: {metrics.box.map:.3f}")
        
    except Exception as e:
        print(f"\\n❌ 训练过程中出现错误: {e}")
        print("\\n故障排除建议:")
        print("1. 检查数据集路径是否正确")
        print("2. 确认data.yaml配置文件存在")
        print("3. 如果是显存不足，减小batch_size")
        print("4. 运行 python check_gpu_support.py 检查GPU状态")

if __name__ == "__main__":
    main()
'''
    
    with open('quick_train_gpu_optimized.py', 'w', encoding='utf-8') as f:
        f.write(fixed_code)
    
    print("✅ 已生成优化的训练脚本: quick_train_gpu_optimized.py")
    print("\n使用方法:")
    print("python quick_train_gpu_optimized.py")

def main():
    """主函数"""
    print("GPU支持诊断工具")
    print("适用于解决PyTorch/YOLO GPU检测问题")
    
    check_system_info()
    check_gpu_hardware()
    check_pytorch_cuda()
    test_gpu_performance()
    provide_solutions()
    generate_quick_train_gpu_fix()
    
    print("\n" + "="*60)
    print("诊断完成!")
    print("="*60)
    
    if torch.cuda.is_available():
        print("🎉 你的系统支持GPU训练!")
        print("建议使用: python quick_train_gpu_optimized.py")
    else:
        print("⚠️ 需要解决GPU支持问题")
        print("请按照上述解决方案进行设置")

if __name__ == "__main__":
    main() 